% 清空环境
% close all; 
clear; clc;
% rng('default'); % 保证可复现性

% =============== 1. 系统参数初始化（表2） ===============
% 阵列参数
Mx = 20; My = 20;          % URA尺寸
M = Mx * My;               % 天线总数
d = 0.5;                   % 半波长间距

% 用户参数
G = 6;                     % 用户组数
Kg = 3;                    % 每组用户数
K = G * Kg;                % 总用户数
P = 10;                    % 每条路径数

% 角度参数（表2）
psi_g = linspace(0, 360, G+1); psi_g(end) = []; % 组平均方位角
theta_g = 73 * ones(1, G); % 组平均仰角
delta_psi = 9;             % 方位角扩展
delta_theta = 12.5;        % 仰角扩展

% 仿真参数
SNR_dB = -40:10:40;          % SNR范围
N_monte = 1000;            % 蒙特卡洛仿真次数

% =============== 2. 生成用户角度信息 ===============
angles = generate_user_angles(K, P, psi_g, theta_g, delta_psi, delta_theta);

% =============== 3. 用户分组（Algorithm 1） ===============
groups = user_grouping(angles);
G_actual = length(unique(groups)); % 实际组数

% =============== 4. 生成信道矩阵 ===============
% 生成所有用户信道
H = generate_channel_matrix(Mx, My, d, P, K, angles);

% =============== 5. 提取AoD支持区域（公式13） ===============
AoD_groups = extract_AoD_support(angles, groups, G_actual);

% =============== 6. RF波束成形设计（Algorithm 2） ===============
F = RF_beamformer_design(Mx, My, d, AoD_groups);
b = size(F, 2); % RF链数量

% =============== 7. 有效信道计算 ===============
H_eff = H' * F; % 降维有效信道

% =============== 8. 混合预编码方案选择 ===============
scheme = 'JGP'; % 可选: 'JGP', 'PGP', 'CGP'

% =============== 9. 主仿真循环 ===============
sum_rate = zeros(length(SNR_dB), 1);

for snr_idx = 1:length(SNR_dB)
    SNR = 10^(SNR_dB(snr_idx)/10);
    P_T = K * SNR; % 总发射功率
    sigma2 = 1;    % 噪声功率归一化

    rate = zeros(N_monte, 1);

    for iter = 1:N_monte % 并行加速
        % 10. 基带预编码设计
        switch scheme
            case 'JGP'
                B = JGP_precoder(H_eff, P_T, sigma2);
            case 'PGP'
                B = PGP_precoder(H_eff, groups, P_T, sigma2);
            case 'CGP'
                B = CGP_precoder(H_eff, AoD_groups, groups, P_T, sigma2);
        end

        % 11. 传输信号生成
        d_symbols = (randn(K, 1) + 1i*randn(K, 1))/sqrt(2); % 数据符号
        s = F * B * d_symbols; % 传输信号

        % 12. 接收信号计算
        r = H' * s + sqrt(sigma2)*(randn(K,1) + 1i*randn(K,1))/sqrt(2);

        % 13. 和速率计算
        SINR_approx = generate_SINR_approx(scheme, H_eff, P_T, sigma2);
        rate(iter) = calculate_sum_rate(H, s, scheme, SINR_approx, sigma2);
    end

    sum_rate(snr_idx) = mean(rate);
end

% =============== 14. 转移块设计（可选） ===============
use_transfer_block = true;
if use_transfer_block
    [T_A, T_B, B_rs] = transfer_block_design(B);
    % 更新传输信号模型: 
    % s = F * (T_A + T_B) * B_rs * d_symbols
end

% =============== 15. 结果可视化 ===============
% 频谱效率曲线
figure;
plot(SNR_dB, sum_rate, 'bo-', 'LineWidth', 2);
xlabel('SNR (dB)');
ylabel('Sum Rate (bps/Hz)');
title(sprintf('3D Hybrid Precoding Performance (%s)', scheme));
grid on;

% 用户分组可视化
% visualize_user_groups(angles, groups);

% AoD支持区域可视化
% visualize_AoD_support(AoD_groups);

fprintf('===== 仿真完成 =====\n');
fprintf('系统配置: %d x %d URA, %d 用户, %d 组\n', Mx, My, K, G_actual);
fprintf('RF链数量: %d (原始), %d (转移块优化后)\n', b, size(B_rs,1));